import os
import csv
import arrow
from app.db.models.tasks import IgnoreWarning, Task
from app.db.models.users import User
from app.db.models.wass import ScanEvents
from app.libs.display import SEVERITY_LEVEL_MAP, IGNORE_RISK_REASON
from app.libs.enums import TASK_TYPE_CONFIG
from app.config.settings import DATETIME_FMT, BASE_DIR


def export_ignore_warnings(username, path=None):
    user = User.objects.filter(username=username).first()
    if not user:
        print(f"{username}不存在")
        exit(-1)
    warnings = IgnoreWarning.objects.find({"uid": user.id, "reason.type": 2})
    if not warnings:
        print(f"{username}没有误报威胁")
        exit(-1)
    print(f"误报威胁{warnings.count()}个")
    task_ids = {}
    for w in warnings:
        task_ids.setdefault(w.taskId, set())
        task_ids[w.taskId].add(w)
    data = []
    for t in Task.objects.filter(id__in=list(task_ids.keys())).only("taskSessionId", "id"):
        data.extend(notify_wecom_ignore_warning(user, t.taskSessionId, list(task_ids.get(t.id))))

    headers = ["用户名", "资产地址", "任务类型", "忽略原因", "忽略结果ID", "忽略时间",
               "威胁地址", "威胁名称", "威胁等级", "task_session_id", "event_id"]
    file = path or os.path.join(BASE_DIR, "ignore_warnings.csv")
    with open(file, "w") as f:
        csv_writer = csv.writer(f)
        csv_writer.writerow(headers)
        csv_writer.writerows(data)


def notify_wecom_ignore_warning(user, task_session_id, warnings):
    data = []
    scan_event_db = ScanEvents._get_collection()
    for w in warnings:
        item = [
            user.username,
            w.targetUrl,
            {TASK_TYPE_CONFIG.get(w.taskType, {}).get('cn')},
            IGNORE_RISK_REASON.get(w.reason.type),
            str(w.id),
            arrow.get(w.ignoredTime).to('local').datetime.strftime(DATETIME_FMT),
            w.detail.affects,
            w.detail.title,
            SEVERITY_LEVEL_MAP[w.detail.level] if w.detail.level >= 2 else '低危',
            task_session_id,
            ""
        ]
        base_q = {"task_session_id": task_session_id}
        q = {}
        if w.taskType == "ssl":
            q = {"detail.name.zh_cn": w.detail.title, **base_q}
        elif w.taskType == "vul":
            q = {
                "detail.name.zh_cn": w.detail.title,
                "$or": [{"detail.affect": w.detail.affects}, {"detail.url": w.detail.affects}],
                **base_q
            }
        elif w.detail.category == "foreign_links":
            q = {"event_name": "risk_link", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "broken_links":
            q = {"event_name": "statistics", "detail.broken_links.referer": w.detail.affects, **base_q}
        elif w.detail.category == "black_links":
            q = {"event_name": "black_links", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "cryjack":
            q = {"event_name": "cryptojacking", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "malscan":
            q = {"event_name": "malscan", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "keyword":
            q = {"event_name": "keyword", "detail.url": w.detail.affects, **base_q}
        elif w.detail.category == "privacy_disclosure":
            q = {"event_name": "privacy_disclosure", "detail.url": w.detail.affects, **base_q}
        if q and (event_ids := [e["event_id"] for e in scan_event_db.find(q)]):
            item[-1] = '、'.join(event_ids)
        data.append(item)
    return data
